const axios = require('axios');
const stream = require('stream');
const { promisify } = require('util');

// Define the necessary variables
const model_name = "THUDM/chatglm3-6b";
// const model_name = "THUDM/LongAlign-13B-64k";
const worker_address = "http://192.168.6.137:17890";
const temperature = 0.0;
const max_new_tokens = 32;
const message = '<|user|>\n你好<|assistant|>';

async function sendTestMessage() {
  try {
    const headers = {
      'User-Agent': 'FastChat Client',
      'Content-Type': 'application/json',
    };
    
    let params = {
      model: model_name,
      prompt: message,
      temperature: temperature,
      max_new_tokens: max_new_tokens,
      stop_token_ids: [64795, 64797, 2],
      echo: false
    };
  
    let pipeline = promisify(stream.pipeline);
  
    let responseStream = await axios({
      method: 'post',
      url: `${worker_address}/worker_generate_stream`,
      headers: headers,
      data: params,
      responseType: 'stream'
    });
  
    let fullResponse = '';
  
    let outStream = new stream.Writable({
      write(chunk, encoding, callback) {
        let part = chunk.toString();
          // 立即将字符串中的Unicode编码转换为字符
        part = part.replace(/\0/g, '');
        part = part.replace(/\\u[\dA-F]{4}/gi, (match) => {
            return String.fromCharCode(parseInt(match.replace(/\\u/g, ''), 16));
        }); 
        fullResponse += part;
        let parsed = JSON.parse(part);
        let text = parsed.text;
        console.clear();
        console.log('Received part of the response:', text);
        callback();
      }
    });
  
    await pipeline(responseStream.data, outStream);
    // console.log('Full response received:', fullResponse);
    
  } catch (error) {
    console.error(`Error sending test message: ${error}`);
  }
}

sendTestMessage();